#include "ransacimagealignment.h"

#include <stdio.h>
#include <iostream>
#include <fstream>

#include "opencv2/core/core.hpp"
#include "opencv2/core/core_c.h"
#include "opencv2/calib3d/calib3d.hpp"

using namespace std;

/******************************************************/
void printMatrix(const cv::Mat& m, const char* name)
{
	cout << name << endl;
	for (int r=0; r<m.rows; ++r)
	{
		for (int c=0; c < m.cols; ++c)
		{
			cout << m.at<double>(r,c) << " ";
		}
		cout << endl;
	}
	cout << endl;
}

/******************************************************/
RansacImageAlignment::RansacImageAlignment():
RansacBase(4,9) // 4 points need to compute solution, 9 doubles for homography matrix solution
{}

/******************************************************/
RansacImageAlignment::~RansacImageAlignment()
{}

/******************************************************/
std::vector<double> RansacImageAlignment::computeSolution(
	std::vector<std::pair<Pt2D, Pt2D> >& samples)
{
	// We solve for the homography mapping using the DLT
	// method. This requires solving the linear system
	// Ax=b.

	// First normalise the points
	std::vector<std::pair<Pt2D, Pt2D> > samples_norm;
	cv::Mat T, T_dash;
	nomalisePoints(samples, samples_norm,T,T_dash);

	// Construct A matrix according to:
	// A_i = [x y 1 0 0 0 ux uy
	//	      0 0 0 x y 1 vx vy] i=1..8
	// b_i = [-u; -v];
	cv::Mat A(8,8,CV_64FC1, 0.0);
	cv::Mat b(8,1,CV_64FC1);
	for (int i=0; i < 4; ++i)
	{
		A.at<double>(2*i,0) = -samples_norm[i].first.x;
		A.at<double>(2*i,1) = -samples_norm[i].first.y;
		A.at<double>(2*i,2) = -1.0;
		A.at<double>(2*i,6) = samples_norm[i].second.x * samples_norm[i].first.x;
		A.at<double>(2*i,7) = samples_norm[i].second.x * samples_norm[i].first.y;
		b.at<double>(2*i,0) = -samples_norm[i].second.x;

		A.at<double>(2*i+1,3) = -samples_norm[i].first.x;
		A.at<double>(2*i+1,4) = -samples_norm[i].first.y;
		A.at<double>(2*i+1,5) = -1.0;
		A.at<double>(2*i+1,6) = samples_norm[i].second.y * samples_norm[i].first.x;
		A.at<double>(2*i+1,7) = samples_norm[i].second.y * samples_norm[i].first.y;
		b.at<double>(2*i+1,0) = -samples_norm[i].second.y;
	}

	//printMatrix(A,"A");

	// Homography H is the solution to Ax=b 
	cv::Mat h_dash;
	cv::solve(A,b,h_dash);

	// Test solution
	cv::Mat H_norm(3,3,CV_64FC1, 0.0);
	H_norm.at<double>(0,0) = h_dash.at<double>(0);
	H_norm.at<double>(0,1) = h_dash.at<double>(1);
	H_norm.at<double>(0,2) = h_dash.at<double>(2);
	H_norm.at<double>(1,0) = h_dash.at<double>(3);
	H_norm.at<double>(1,1) = h_dash.at<double>(4);
	H_norm.at<double>(1,2) = h_dash.at<double>(5);
	H_norm.at<double>(2,0) = h_dash.at<double>(6);
	H_norm.at<double>(2,1) = h_dash.at<double>(7);
	H_norm.at<double>(2,2) = 1.0;

	//printMatrix(H_norm, "H_norm");

	cv::Mat T_dash_inv = T_dash.inv();
	cv::Mat H = T_dash_inv * H_norm * T;

	vector<double> h(9); 
	h[0] = H.at<double>(0,0)/H.at<double>(2,2);
	h[1] = H.at<double>(0,1)/H.at<double>(2,2);
	h[2] = H.at<double>(0,2)/H.at<double>(2,2);
	h[3] = H.at<double>(1,0)/H.at<double>(2,2);
	h[4] = H.at<double>(1,1)/H.at<double>(2,2);
	h[5] = H.at<double>(1,2)/H.at<double>(2,2);
	h[6] = H.at<double>(2,0)/H.at<double>(2,2);
	h[7] = H.at<double>(2,1)/H.at<double>(2,2);
	h[8] = H.at<double>(2,2)/H.at<double>(2,2);

	//printMatrix(T, "T");
	//printMatrix(T_dash, "T_dash");
	//printMatrix(T_dash_inv, "T_dash_inv");
	//printMatrix(H, "H");
	//
	//cv::Mat u(3,1,CV_64FC1,1.0);
	//cv::Mat v = H*u;
	//
	//printMatrix(v, "v");

	//std::vector<cv::Point2d> src(4), dst(4);
	//for (int i=0; i < 4; ++i)
	//{
	//	src[i].x = samples[i].first.x;
	//	src[i].y = samples[i].first.y;
	//	dst[i].x = samples[i].second.x;
	//	dst[i].y = samples[i].second.y;
	//}
	//
	//cv::Mat H_cv = cv::findHomography(cv::Mat(src), cv::Mat(dst), CV_RANSAC);
	//printMatrix(H_cv, "H_cv");
	//
	return h;
}

/******************************************************/
void RansacImageAlignment::nomalisePoints(
	const std::vector<std::pair<Pt2D, Pt2D> >& samples,
	std::vector<std::pair<Pt2D, Pt2D> >& samples_normalised,
	cv::Mat& T,
	cv::Mat& T_dash)
{
	// Compute the normalising transformations T and T_dash
	std::vector<Pt2D> points_src(samples.size());
	std::vector<Pt2D> points_dst(samples.size());
	for (unsigned int i=0; i < samples.size(); ++i)
	{
		points_src[i] = samples[i].first;
		points_dst[i] = samples[i].second;
	}

	normalise(points_src, T);
	normalise(points_dst, T_dash);

	// Normalise the source points with T and T_dash
	samples_normalised.resize(samples.size());
	for (unsigned int i=0; i < samples.size(); ++i)
	{
		cv::Mat x(3,1,CV_64FC1,1.0);
		cv::Mat x_dash(3,1,CV_64FC1,1.0);

		x.at<double>(0) = samples[i].first.x;
		x.at<double>(1) = samples[i].first.y;

		x_dash.at<double>(0) = samples[i].second.x;
		x_dash.at<double>(1) = samples[i].second.y;

		cv::Mat x_norm = T * x;
		cv::Mat x_dash_norm = T_dash * x_dash;

		samples_normalised[i].first.x = x_norm.at<double>(0);
		samples_normalised[i].first.y = x_norm.at<double>(1);

		samples_normalised[i].second.x = x_dash_norm.at<double>(0);
		samples_normalised[i].second.y = x_dash_norm.at<double>(1);
	}
}

/******************************************************/
void RansacImageAlignment::normalise(
	const std::vector<Pt2D>& points,
	cv::Mat& T)
{
	assert(points.size() > 0);

	// Compute offset to center points on origin
	double y_avg = 0.0;
	double x_avg = 0.0;
	for (unsigned int i = 0; i < points.size(); ++i)
	{
		x_avg += points[i].x;
		y_avg += points[i].y;
	}
	y_avg /= points.size();
	x_avg /= points.size();

	// Compute scale to have avergage dist from origin sqrt(2)
	double dist_avg = 0;
	for (unsigned int i = 0; i < points.size(); ++i)
	{
		const double a = points[i].x - x_avg;
		const double b = points[i].y - y_avg;
		dist_avg += sqrt(a*a + b*b);
	}
	dist_avg /= points.size();
	const double scale = sqrt(2.0)/dist_avg;

	T = cv::Mat(3,3,CV_64FC1,0.0);
	T.at<double>(0,0) = scale;
	T.at<double>(1,1) = scale;
	T.at<double>(2,2) = 1.0;
	T.at<double>(0,2) = -x_avg*scale;
	T.at<double>(1,2) = -y_avg*scale;
}

/******************************************************/
double RansacImageAlignment::evaluateSolution(
	std::vector<double>& solution, 
	std::vector<int>& inliers)
{
	inliers.clear();

	cv::Mat H(3,3,CV_64FC1);
	H.at<double>(0,0) = solution[0];
	H.at<double>(0,1) = solution[1];
	H.at<double>(0,2) = solution[2];
	H.at<double>(1,0) = solution[3];
	H.at<double>(1,1) = solution[4];
	H.at<double>(1,2) = solution[5];
	H.at<double>(2,0) = solution[6];
	H.at<double>(2,1) = solution[7];
	H.at<double>(2,2) = solution[8];

	for (unsigned int i=0; i < _samples.size(); ++i)
	{
		cv::Mat u(3,1,CV_64FC1);
		u.at<double>(0,0) = _samples[i].first.x;
		u.at<double>(1,0) = _samples[i].first.y;
		u.at<double>(2,0) = 1.0;

		cv::Mat v = H*u;

		double x = u.at<double>(0,0)/u.at<double>(2,0);
		double y = u.at<double>(1,0)/u.at<double>(2,0);

		double xt = v.at<double>(0,0)/v.at<double>(2,0);
		double yt = v.at<double>(1,0)/v.at<double>(2,0);

		if (v.at<double>(2,0) != 0.0)
		{
			const double a = v.at<double>(0,0)/v.at<double>(2,0) - _samples[i].second.x;
			const double b = v.at<double>(1,0)/v.at<double>(2,0) - _samples[i].second.y;
			const double dist = sqrt(a*a + b*b);
			if (dist < _inlier_threshold)
			{
				inliers.push_back(i);
			}
		}

	}
	return (double)inliers.size();
}
